import os
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
import random

class RibFracClassifierDataset(Dataset):
    def __init__(self, data_dir, mode='train', transform=None):
        self.data_dir = data_dir
        self.mode = mode
        self.transform = transform
        
        metadata_path = os.path.join(data_dir, 'metadata', 'metadata.csv')
        if not os.path.exists(metadata_path):
            #进入推理模式
            self.metadata = None
            print(f"元数据文件不存在: {metadata_path}；进入推理模式")
            # raise FileNotFoundError(f"元数据文件不存在: {metadata_path}")
        else:
            self.metadata = pd.read_csv(metadata_path)
            # self.metadata['class_label'] = self.metadata['class_label'].replace(-1, 0) # 保留-1标签为0
            self.metadata = self.metadata[self.metadata['class_label'] != -1] #删除标签-1
            self.metadata = self.metadata[self.metadata['class_label'] < 5]  # 确保只有0-4
            self.metadata['class_label'] = self.metadata['class_label'] # 0-based
            self.metadata['identifier'] = self.metadata['public_id'] + '_' + self.metadata['label_id'].astype(str)
         
            self.metadata.reset_index(drop=True, inplace=True)# 舍弃掉不使用的标签，之后进行重新排列
            print(f"加载{mode}数据集: {len(self.metadata)}个样本")
            print(f"骨折类型分布:")
            print(self.metadata['class_label'].value_counts())
            self.class_count = self.metadata['class_label'].value_counts()

    def __class_counts__(self):
        return self.class_count
    
    def __len__(self):
        return len(self.metadata)
    
    def get_labels(self):
        return self.metadata['class_label'].values
    
    def get_class_indices(self):
        class_indices = {i: [] for i in range(5)}  # 0-4类
        for idx, label in enumerate(self.get_labels()):
            class_indices[label].append(idx)
        return class_indices
    
    def __getitem__(self, idx):
        row = self.metadata.iloc[idx]
        ct_path = os.path.join(self.data_dir, 'images', os.path.basename(row['ct_path']))
        if not os.path.exists(ct_path):
            raise FileNotFoundError(f"CT文件不存在: {ct_path}")
        ct_block = np.load(ct_path)  # (64, 64, 64)
        start_idx = (64 - 32) // 2
        cropped_block = ct_block[
            start_idx:start_idx+32, 
            start_idx:start_idx+32,
            start_idx:start_idx+32
        ]
        ct_block = cropped_block[np.newaxis, ...]  # (1, 64, 64, 64)
        class_label = row['class_label']  # 已经是0-based (0-3)
        
        min_val = np.min(ct_block)
        max_val = np.max(ct_block)
        
        if max_val - min_val <= 1e-6: 
            if max_val != 0:
                ct_block_normalized = ct_block / max_val  # 避免除以0
            else:
                ct_block_normalized = np.zeros_like(ct_block)
        else:
            ct_block_normalized = (ct_block - min_val) / (max_val - min_val)
        
        mean = np.mean(ct_block_normalized)
        std = np.std(ct_block_normalized)
        if std > 1e-6:
            ct_block_normalized = (ct_block_normalized - mean) / std

            max_abs = np.max(np.abs(ct_block_normalized))
            if max_abs > 0:
                ct_block_normalized /= max_abs

        ct_tensor = torch.from_numpy(ct_block_normalized).float()

        class_label = row['class_label']  # 0-based (0-4)
        identifier = row['identifier']

        if self.transform and self.mode == 'train':
            ct_tensor = self.transform(ct_tensor)
        
        return ct_tensor, torch.tensor(class_label, dtype=torch.long), identifier
    
class Random3DFlip:
    """3D随机翻转"""
    def __init__(self, p=0.5):
        self.p = p
        
    def __call__(self, x):
        if random.random() < self.p:
            axis = random.choice([1, 2, 3])  # 1:D, 2:H, 3:W
            return torch.flip(x, dims=[axis])
        return x
    
class Random3DRotation:
    def __init__(self, p=0.3):
        self.p = p
        
    def __call__(self, x):
        if random.random() < self.p:
            k = random.randint(0, 3)
            plane = random.choice([(1, 2), (1, 3), (2, 3)])
            return torch.rot90(x, k, dims=plane)
        return x

def get_data_transforms(mode):
    if mode == 'train':
        return transforms.Compose([
            Random3DRotation(p=0.3),
            Random3DFlip(p=0.3)
        ])
    else:
        return None

import torch
import numpy as np

class Class3EnsuredBatchSampler(torch.utils.data.Sampler):
    def __init__(self, dataset, batch_size):
        self.dataset = dataset
        self.batch_size = batch_size
        self.class_indices = dataset.get_class_indices()
        
        self.num_class4 = len(self.class_indices[4])
        if self.num_class4 == 0:
            raise ValueError("数据集中不存在类别3的样本")
        
        self.length = self.num_class4
        
        self.non4_indices = []
        for c in [0, 1, 2, 3]:
            self.non4_indices.extend(self.class_indices[c])
        self.non4_count = len(self.non4_indices)
        
        np.random.shuffle(self.non4_indices)
        self.non3_index = 0
        
    def __iter__(self):
        class4_indices = self.class_indices[4].copy()
        np.random.shuffle(class4_indices)
        
        for class4_idx in class4_indices:
            batch = [class4_idx]
            
            for _ in range(self.batch_size - 1):
                if self.non4_index >= self.non4_count:
                    np.random.shuffle(self.non4_indices)
                    self.non4_index = 0
                
                batch.append(self.non4_indices[self.non4_index])
                self.non4_index += 1
                
            yield batch
    
    def __len__(self):
        return self.length

def create_classifier_loaders(data_root, batch_size=16):
    train_dir = os.path.join(data_root, 'train')
    val_dir = os.path.join(data_root, 'val')
    
    train_dataset = RibFracClassifierDataset(
        train_dir,
        mode='train',
        transform = get_data_transforms('train')
    )
    val_dataset = RibFracClassifierDataset(
        val_dir,
        mode='val'
    )

    # batch_sampler = Class3EnsuredBatchSampler(train_dataset, batch_size=16)

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        # batch_sampler=batch_sampler,
        pin_memory=True
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=True,
        pin_memory=True
    )

    
    return train_loader, val_loader

def create_test_classifier_loaders(data_root, batch_size=16):
    test_dir = os.path.join(data_root)
    test_dataset = RibFracClassifierDataset(
        test_dir,
        mode='test'
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        pin_memory=True
    )
    return test_loader

if __name__ == "__main__":
    data_root = './ribfrac-dataset/z_spacing'
    
    print("创建数据加载器...")
    train_loader, val_loader = create_classifier_loaders(
        data_root,
        batch_size=4
    )
    print("\n可视化训练批次...")
